#!/usr/bin/env python3
import sqlite3
import zipfile
from pathlib import Path
from typing import List, Dict, Optional
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn
from rich.traceback import install
from rich.logging import RichHandler
from rich.table import Table
from rich import print
import logging
import argparse

console = Console()

# 日志配置

install(show_locals=True)

# 配置 Rich 日志处理器（替换原 basicConfig）
logging.basicConfig(
    level=logging.INFO,
    format="%(message)s",
    handlers=[RichHandler(
        rich_tracebacks=True,          # 启用彩色追踪
        markup=True,                   # 允许日志内容使用 Markdown
        show_time=True,               # 时间由 rich 自动处理
        show_level=True,               # 显示日志级别图标
        show_path=False                # 简化文件路径显示
    )]
)

# 日志配置

# 类型映射配置
TYPE_MAPPING = {
    "INTEGER": "INT",
    "REAL": "FLOAT",
    "TEXT": "TEXT",
    "BLOB": "BLOB",
    "NUMERIC": "NUMERIC",
}

class DatabaseExporter:
    def __init__(self, db_path: str):
        self.db_path = Path(db_path)
        self.conn = sqlite3.connect(str(self.db_path))
        self.cursor = self.conn.cursor()
        
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.conn.close()

    def get_all_tables(self) -> List[str]:
        """获取所有用户表"""
        self.cursor.execute(
            "SELECT name FROM sqlite_master "
            "WHERE type='table' AND name NOT LIKE 'sqlite_%'"
        )
        return [row[0] for row in self.cursor.fetchall()]

    def get_table_schema(self, table_name: str) -> Dict:
        """获取表结构信息"""
        self.cursor.execute(f"PRAGMA table_info({table_name})")
        columns = []
        for row in self.cursor.fetchall():
            col_type = row[2].upper().split("(")[0]
            columns.append({
                "name": row[1],
                "type": TYPE_MAPPING.get(col_type, "TEXT"),
                "nullable": not row[3],
                "default": row[4]
            })
        return {"columns": columns}

    def generate_create_table(self, table_name: str) -> str:
        """生成CREATE TABLE语句"""
        schema = self.get_table_schema(table_name)
        columns = []
        for col in schema["columns"]:
            col_def = f"{col['name']} {col['type']}"
            if not col["nullable"]:
                col_def += " NOT NULL"
            if col["default"]:
                col_def += f" DEFAULT {col['default']}"
            columns.append(col_def)
        return f"CREATE TABLE {table_name} (\n  " + ",\n  ".join(columns) + "\n);\n"

    def export_table_data(self, table_name: str) -> str:
        """导出表数据为INSERT语句"""
        insert_sql = []
        self.cursor.execute(f"SELECT * FROM {table_name}")
        columns = [desc[0] for desc in self.cursor.description]
        
        for row in self.cursor.fetchall():
            values = []
            for value in row:
                if value is None:
                    values.append("NULL")
                elif isinstance(value, (int, float)):
                    values.append(str(value))
                else:
                    escaped = str(value).replace("'", "''")
                    values.append(f"'{escaped}'")
            insert_sql.append(
                f"INSERT INTO {table_name} ({', '.join(columns)}) "
                f"VALUES ({', '.join(values)});"
            )
        return "\n".join(insert_sql) + "\n"

def main(db_path: str, zip_path: str):
    """主处理流程"""
    zip_path = Path(zip_path)
    sql_file = zip_path.with_suffix(".sql")
    
    try:
        with DatabaseExporter(db_path) as exporter, \
             Progress(SpinnerColumn(), *Progress.get_default_columns(), TimeElapsedColumn()) as progress:
            
            # 获取所有表
            task1 = progress.add_task("[cyan]正在读取数据库结构...", total=1)
            tables = exporter.get_all_tables()
            progress.update(task1, completed=1)
            
            if not tables:
                console.print("[red]错误：数据库中没有用户表！")
                return

            # 显示表信息
            table_info = Table("表名", "列数", "预估数据量", title="数据库结构分析")
            for table in tables:
                exporter.cursor.execute(f"SELECT COUNT(*) FROM {table}")
                count = exporter.cursor.fetchone()[0]
                table_info.add_row(
                    table,
                    str(len(exporter.get_table_schema(table)["columns"])),
                    f"{count:,} 行"
                )
            console.print(table_info)

            # 生成SQL文件
            task2 = progress.add_task("[green]生成SQL文件...", total=len(tables)*2)
            with open(sql_file, "w", encoding="utf-8") as f:
                for table in tables:
                    # 生成表结构
                    f.write(exporter.generate_create_table(table) + "\n")
                    progress.update(task2, advance=1)
                    
                    # 导出数据
                    data = exporter.export_table_data(table)
                    f.write(data + "\n")
                    progress.update(task2, advance=1)
                    progress.log(f"[yellow]已导出表：{table}")

            # 创建ZIP压缩包
            task3 = progress.add_task("[blue]创建压缩文件...", total=1)
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                zipf.write(sql_file, arcname=sql_file.name)
            progress.update(task3, completed=1)

        # 清理临时文件
        sql_file.unlink()
        console.print(f"[bold green]\n✅ 成功生成压缩文件：{zip_path}")

    except Exception as e:
        console.print(f"[red bold]\n⚠️ 处理过程中发生错误：{str(e)}")
        if sql_file.exists():
            console.print(f"[yellow]临时文件保留在：{sql_file}")
        
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="📦 SQLite数据库转换工具",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("-d", "--database", required=True, 
                      help="输入SQLite数据库文件路径")
    parser.add_argument("-z", "--zip", required=True,
                      help="输出ZIP文件路径")
    
    args = parser.parse_args()
    
    console.print("[bold]🔧 SQLite数据库转换工具[/bold]", justify="center")
    console.print(f"输入数据库：{args.database}")
    console.print(f"输出位置：{args.zip}\n")
    
    main(args.database, args.zip)
